#!/usr/bin/python
# GM_SimpleIntegrate.py

#import math
import numpy
from time import asctime
from os import linesep

# Must have G, c in appropriate units so rE is in pc! Note M is in Msun
# See notes pg. 13 date 2/22/11
Goverc2 = 4.786e-14

VariableNames = ['appearance', 'masstuple', 'vtuple', 'd', 'deltad', 'ba', 'deltaba', 'deltala', 'MvVals', 'delMvVals', 'MvMass', 'delMvMass', 'dlocs', 'blocs', 'SourceMag', 'delMvList', 'darkstars']


def GetEinsteinModifier(Mv):
    return 1.


def simple(simpars, phypars, appearance, masstuple, vtuple):
    '''Acquire lensing rates, qualified only by source magnitude.'''
    # Intake 'n' stuff
    for i in simpars:
        cmd = "%s = simpars['%s']" % (i,i)
        exec cmd
    for i in phypars:
        cmd = "%s = phypars['%s']" % (i,i)
        exec cmd
    if len(appearance) == 1:
        print "Process", codeindex, "beginning."
        print asctime()
    (darkstars, rho0mlist) = masstuple
    (vrelappl, sigmaappl2, sigmaappb2) = vtuple
    # Distance elements
    DL = numpy.array(d)
    dDL = numpy.array(deltad)
    DLdDL = DL*dDL
    DS = numpy.array(d)
    dDS = numpy.array(deltad)
    DS2dDS = DS*DS*dDS
    # Angular element, incl. symmetry restoration via quadrupling
    cosb = numpy.cos(ba)
    cosbdeltab = cosb*deltaba
    dOmega = numpy.multiply(deltala[:,numpy.newaxis], cosbdeltab[numpy.newaxis,:])*4 # Symmetry!
    del cosb, cosbdeltab
    # These histograms compile things already "naturally" separated and hold a running tally at any point in the program.  DSHist, DLHist, and bHist are collected before final summation for each 
    MvHist = numpy.zeros(len(delMvVals), dtype=float) # Divide by dMv here
    MassHist = numpy.zeros(len(MvMass), dtype=float) # Divide by dM here
    DSHist = numpy.zeros(len(dlocs), dtype=float) # Divide by dDS here
    DLHist = numpy.zeros(len(dlocs), dtype=float) # Divide by dDL here
    bHist = numpy.zeros(len(blocs), dtype=float) # Divide by dba here
    # These "master histograms" contain snapshots of the above histograms at each of the limiting source magnitude cutoffs.
    MasterMvHist = []
    MasterMassHist = []
    MasterDSHist = []
    MasterDLHist = []
    MasterbHist = []
    # These histograms need additional complex masks to delimit them
#    vtappHist =  numpy.ones(110, dtype=float) #### Not supported yet
    # This holds the overall result
    lensrate = 0.0
    lensrates = numpy.zeros(len(SourceMag))
    SourceMagIndex = 0
    CurrentLimitingMag = SourceMag[SourceMagIndex]
    # i indexes source apparent magnitude
    for i in range(0, len(appearance)):
        nS = numpy.array(appearance[i])
        Mv = float(MvVals[i]) # Intrinsic magnitudes are from MvList, but apparent magnitudes are from MvVals
        dMv = float(delMvList[i])
        EinsteinModifier = GetEinsteinModifier(Mv)
        # Cut for source unlensed magnitude
        if Mv > (CurrentLimitingMag + 0.00001):
            lensrates[SourceMagIndex] = lensrate
            MasterMvHist.append(numpy.array(MvHist))
            MasterMassHist.append(numpy.array(MassHist))
            MasterDSHist.append(numpy.array(DSHist))
            MasterDLHist.append(numpy.array(DLHist))
###
            # If you've reached the end, stop; otherwise, move to next stage
            if SourceMagIndex == (len(SourceMag) - 1):
                break
            else:
                SourceMagIndex = SourceMagIndex + 1
                CurrentLimitingMag = SourceMag[SourceMagIndex]
###
        icontrib = 0.0
        # j indexes lens mass
        for j in range(0, len(darkstars)):
            nL = darkstars[j]
            M = MvMass[j,1]
            ijcontrib = 0.0
            # k indexes distance or position
            for k in range(0, len(dlocs)):
                Ki = k
                Kf = len(dlocs) - k
                # 1d
                # two_rE:  2x Einstein radius, times a modifier for photometric sensitivity, if available
                # DLdDL_DS2dDS:  Volume elements & additional distance terms to correct
                # AllDistances:  Combine two_rE and DLdDL_DS2dDS
                two_rE = 2*(4*Goverc2*M*DL[:Kf]*(DS[Ki:] - DL[:Kf])/DS[Ki:])**0.5*EinsteinModifier
                DLdDL_DS2dDS = DLdDL[:Kf]*DS2dDS[Ki:]
                AllDistances = two_rE*DLdDL_DS2dDS
                # 3d
                # RMSvtapp:  RMS angular velocity difference
                # VolumeElements:  Combine angular-only and distance-only terms
                # GammaK:  Combine all.  GammaK gives the lensing contribution at each location for this staggering of this lens and source population
                RMSvtapp = ((vrelappl[Ki:,:,:] - vrelappl[:Kf,:,:])**2 + sigmaappl2[Ki:,:,:] + sigmaappl2[:Kf,:,:] + sigmaappb2[Ki:,:,:] + sigmaappb2[:Kf,:,:])**0.5
                VolumeElements = numpy.multiply(AllDistances[:,numpy.newaxis,numpy.newaxis], dOmega[numpy.newaxis,:,:])
                GammaK = nS[Ki:,:,:]*nL[:Kf,:,:]*RMSvtapp*VolumeElements
                # Sum over l, then d, to have b-dependence
                GammaKldsummed = (GammaK.sum(axis=1)).sum(axis=0)
                bHist = bHist + GammaKldsummed
                # Sum over b, then l, to have D-dependence
                GammaKblsummed = (GammaK.sum(axis=2)).sum(axis=1)
                DSHist[Ki:] = DSHist[Ki:] + GammaKblsummed
                DLHist[:Kf] = DLHist[:Kf] + GammaKblsummed
                # Final summation to get total rate from this staggering/source/lens integrand term
                ijkcontrib = GammaKblsummed.sum()
                ijcontrib = ijcontrib + ijkcontrib
            MassHist[j] = MassHist[j] + ijcontrib
            icontrib = icontrib + ijcontrib
        lensrate = lensrate + icontrib
        MasterbHist.append(numpy.array(bHist)) # Unlike others, it is important to collect this quantity for each magnitude
        MvHist[i] = icontrib
        print "For a width of", dMv, "around a source magnitude of", Mv, "I found a lensing rate contribution of", icontrib
        print asctime()
        # One last pass to catch anything that made it all the way through
        if i == (len(appearance) - 1):
            lensrates[SourceMagIndex] = lensrate
            MasterMvHist.append(numpy.array(MvHist))
            MasterMassHist.append(numpy.array(MassHist))
            MasterDSHist.append(numpy.array(DSHist))
            MasterDLHist.append(numpy.array(DLHist))
#            MasterbHist.append(numpy.array(bHist))
    if len(appearance) != 1:
        for i in range(0, len(SourceMag)):
            print "For a limiting unlensed source magnitude of", str(SourceMag[i])+ ", a lensing rate of", str(lensrates[i]), "was found."
    if len(appearance) == 1:
        print "Process", codeindex, "finishing."
        print asctime()
    # Divide each histogram by the width of the bin (i.e., normalize).
    for i in range(0, len(MasterMvHist)):
        MasterMvHist[i] = MasterMvHist[i]/delMvVals[:]
        MasterMassHist[i] = MasterMassHist[i]/delMvMass[:,1]
        MasterDSHist[i] = MasterDSHist[i]/dDS
        MasterDLHist[i] = MasterDLHist[i]/dDL
    for i in range(0, len(MasterbHist)):
        MasterbHist[i] = MasterbHist[i]/deltaba
    print "GM_SimpleIntegrate ran just fine, thx for asking."
    return (lensrates, MasterMvHist, MasterMassHist, MasterDSHist, MasterDLHist, MasterbHist)
                


############################
# COMMANDLINE OPERATION courtesy of codemaster Rachel Street!

if __name__ == '__main__':
    # This provision is for multithreading only, really
    # The comstr will be of the form 'python thismodule parfile outfile'

    # Import required functions:
    from sys import argv
    import cPickle as cP
    
    # Read commandline argument indicating the location of the input arguments
    parfilename = argv[-2]

    # Open and unpack
    parfile = open(parfilename, 'r')
    pars = cP.load(parfile)
    parfile.close()
    for i in pars:
        cmd = "%s = pars['%s']" % (i,i)
        exec cmd

    # Call function
    ratestuple = simple(simpars, phypars, appearance, masstuple, vtuple)

    # Read commandline argument indicating the output location
    outfilename = argv[-1]

    # Pack and write
    outpars = {'simpars' : simpars,
               'phypars' : phypars,
               'ratestuple' : ratestuple}
    outfile = open(outfilename, 'w')
    cP.dump(outpars, outfile)
    outfile.close()
